#include <stdint.h>
#include <stdio.h>

#include "../../inc/dilithium/dilithium.h"
#include "../../inc/dilithium/params.h"
#include "../../inc/dilithium/sign.h"
#include "../../inc/rand/randombytes.h"

#if REE_OR_TEE == 1

#include <tee_ext_api.h>
#include <tee_log.h>
#include <securec.h>

#include "tee_defines.h"

#include "tee_crypto_api.h"
#include "fips202.h"

#define TEE_TIME_MILLIS_BASE    1000

#define TEE_TIME_SUB(t1, t2, dst) do {                      \
        (dst).seconds = (t1).seconds - (t2).seconds;        \
        if ((t1).millis < (t2).millis) {                    \
            (dst).seconds--;                                \
            (dst).millis = (t1).millis + TEE_TIME_MILLIS_BASE - (t2).millis;\
        } else {                                            \
            (dst).millis = (t1).millis - (t2).millis;       \
        }                                                   \
    } while (0)

uint32_t tee_time_sub(TEE_Time * before, TEE_Time* after)
{
    TEE_Time dst;
    TEE_TIME_SUB(*after, *before, dst);
    return dst.seconds * TEE_TIME_MILLIS_BASE + dst.millis;
}

// #include <time.h>

typedef TEE_Time bench_t;

typedef struct _bench_st
{
 /** Stores the time measured before the execution of the benchmark. */
 bench_t before;
 /** Stores the time measured after the execution of the benchmark. */
 bench_t after;
 /** Stores the sum of timings for the current benchmark. */
 uint32_t total;

 double	result;
} bench_ctx;


void bench_ctx_reset(bench_ctx * ctx) {
    ctx->total = 0;
}

void bench_ctx_before(bench_ctx * ctx)
{
    TEE_GetSystemTime(&ctx->before);
}

void bench_ctx_after(bench_ctx * ctx)
{
    uint32_t result;

    TEE_GetSystemTime(&ctx->after);
    // millisecond
    // result = ctx->after.millis - ctx->before.millis;
    ctx->total += tee_time_sub(&ctx->before, &ctx->after);
}

void bench_ctx_compute(bench_ctx * ctx, int benches)
{
    ctx->result = (double)ctx->total / benches;
}


void bench_ctx_print(const bench_ctx * ctx, const char *label)
{
    tloge("BENCH: %s %*c = %9.2f ms, total = %u ms", label, (int)(32 - strlen(label)), ' ', ctx->result, ctx->total);
}

#define MLEN 100

#define ROUND 10

void test_dilithium(int level)
{
    tloge("+------------------------------------------------------------------------------+\n");
    dilithium_init(level);
    switch (level) {
        case DILITHIUM_NIST2:
            tloge("+---    Test module  --->  Dilithium Level %d.\n", 2);
            break;
        case DILITHIUM_NIST3:
            tloge("+---    Test module  --->  Dilithium Level %d.\n", 3);
            break;
        case DILITHIUM_NIST5:
            tloge("+---    Test module  --->  Dilithium Level %d.\n", 5);
            break;
        default:
            tloge("ERROR LEVEL!!!\n");
            break;
    }
    unsigned int i, j;
    int ret;
    size_t mlen, smlen;
    uint8_t m[MLEN] = {0};
    uint8_t sm[MLEN + MAX_CRYPTO_BYTES];
    uint8_t m2[MLEN + MAX_CRYPTO_BYTES];
    uint8_t pk[MAX_CRYPTO_PUBLICKEYBYTES];
    uint8_t sk[MAX_CRYPTO_SECRETKEYBYTES];

    for(i = 0; i < 100; ++i) {
        // tloge("randombytes before");
        randombytes(m, MLEN);
        // tloge("randombytes after");

        crypto_sign_keypair(pk, sk);
        crypto_sign(sm, &smlen, m, MLEN, sk);
        ret = crypto_sign_open(m2, &mlen, sm, smlen, pk);

        if(ret) {
            tloge("Verification failed [Fail!!!!]\n");
            goto END;
        }

        if(mlen != MLEN) {
            tloge("Message lengths don't match [Fail!!!!]\n");
            goto END;
        }

        for(j = 0; j < mlen; ++j) {
            if(m[j] != m2[j]) {
                tloge("Messages don't match [Fail!!!!]\n");
                goto END;
            }
        }

        randombytes((uint8_t *)&j, sizeof(j));
        do {
            randombytes(m2, 1);
        } while(!m2[0]);
        sm[j % params.CRYPTO_BYTES] += m2[0];
        ret = crypto_sign_open(m2, &mlen, sm, smlen, pk);
        if(!ret) {
            tloge("Trivial forgeries possible [Fail!!!!]\n");
            goto END;
        }
    }

    tloge("CRYPTO_PUBLICKEYBYTES = %d\n", params.CRYPTO_PUBLICKEYBYTES);
    tloge("CRYPTO_SECRETKEYBYTES = %d\n", params.CRYPTO_SECRETKEYBYTES);
    tloge("CRYPTO_BYTES = %d\n", params.CRYPTO_BYTES);

    tloge("Dilithium [PASS]\n");

    END:
    tloge("+------------------------------------------------------------------------------+\n");
}

void bench_dilithium(uint32_t out[3], int level)
{
    // tloge("+------------------------------------------------------------------------------+\n");
    dilithium_init(level);
    switch (level) {
        case DILITHIUM_NIST2:
            tloge("+---    Bench module  --->  Dilithium Level %d.\n", 2);
            break;
        case DILITHIUM_NIST3:
            tloge("+---    Bench module  --->  Dilithium Level %d.\n", 3);
            break;
        case DILITHIUM_NIST5:
            tloge("+---    Bench module  --->  Dilithium Level %d.\n", 5);
            break;
        default:
            tloge("ERROR LEVEL!!!\n");
            break;
    }

    unsigned int i, j;
    size_t mlen, smlen;
    uint8_t m[MLEN] = {0};
    uint8_t sm[MLEN + MAX_CRYPTO_BYTES];
    uint8_t m2[MLEN + MAX_CRYPTO_BYTES];
    uint8_t pk[MAX_CRYPTO_PUBLICKEYBYTES];
    uint8_t sk[MAX_CRYPTO_SECRETKEYBYTES];

    bench_ctx keygen, sign, vrf;
    bench_ctx_reset(&keygen);
    bench_ctx_reset(&sign);
    bench_ctx_reset(&vrf);

    for(i = 0; i < NTESTS/ROUND; ++i) {
        randombytes(m, MLEN);

        bench_ctx_before(&keygen);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign_keypair(pk, sk);
        }
        bench_ctx_after(&keygen);

        bench_ctx_before(&sign);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign(sm, &smlen, m, MLEN, sk);
        }
        
        bench_ctx_after(&sign);

        bench_ctx_before(&vrf);
        for (j = 0; j < ROUND; j++)
        {
            crypto_sign_open(m2, &mlen, sm, smlen, pk);
        }
        bench_ctx_after(&vrf);
    }

    bench_ctx_compute(&keygen, NTESTS);
    bench_ctx_compute(&sign, NTESTS);
    bench_ctx_compute(&vrf, NTESTS);

    out[0] = keygen.total;
    out[1] = sign.total;
    out[2] = vrf.total;

    tloge("crypto_sign_keypair = %.3f ms", keygen.result);
    tloge("crypto_sign = %.3f ms", sign.result);
    tloge("crypto_sign_open = %.3f ms", vrf.result);

    // bench_ctx_print(&keygen, "crypto_sign_keypair");
    // bench_ctx_print(&sign, "crypto_sign");
    // bench_ctx_print(&vrf, "crypto_sign_open");

    // tloge("+------------------------------------------------------------------------------+\n");
}

#define BUF_SIZE (32 * 256)
#define INBUF_SIZE 32
void bench_rng(uint32_t out[3]){
    uint8_t inBuf[INBUF_SIZE], outBuf[BUF_SIZE];
    bench_ctx rnd, shk128, shk256;

    bench_ctx_reset(&rnd);
    // bench_ctx_reset(&shk128);
    // bench_ctx_reset(&shk256);
    
    bench_ctx_before(&rnd);
    for (int i = 0; i < NTESTS; i++)
    {
        TEE_GenerateRandom(inBuf, 32);
    }
    bench_ctx_after(&rnd);

    bench_ctx_before(&shk128);
    for (int i = 0; i < NTESTS; i++)
    {
        shake128(outBuf, BUF_SIZE, inBuf, 32);
    }
    bench_ctx_after(&shk128);

    bench_ctx_before(&shk256);
    for (int i = 0; i < NTESTS; i++)
    {
        shake256(outBuf, BUF_SIZE, inBuf, 32);
    }
    bench_ctx_after(&shk256);

    bench_ctx_compute(&rnd, NTESTS);
    bench_ctx_compute(&shk128, NTESTS);
    bench_ctx_compute(&shk256, NTESTS);

    out[0] = rnd.total;
    out[1] = shk128.total;
    out[2] = shk256.total;

    tloge("TEE_GenerateRandom = %.3f ms", rnd.result);
    tloge("shake128 = %.3f ms", shk128.result);
    tloge("shake256 = %.3f ms", shk256.result);
}

#endif // REE_OR_TEE == 1